[EagerJIT] Fix bug in handling of positional arguments by kurisu6912 · Pull Request #1675 · tile-ai/tilelang

kurisu6912 · 2026-01-15T03:19:26Z

This pr fix some wrong code in jit argument parsing, which results in a bug that positional arguments are ignored.

import tilelang
import tilelang.language as T


@tilelang.jit
def vectorize_broadcast_int8(vec_num):
    with T.Kernel(1, threads=128):
        a = T.alloc_local((64,), 'int8')
        b = T.alloc_var('int8')
        for i in T.vectorized(vec_num):
            a[i] = b


kernel = vectorize_broadcast_int8.compile(128)
print(kernel)

Summary by CodeRabbit

Style
- Adjusted formatting in code examples with blank line and quote standardization.
Refactor
- Improved internal argument handling logic for better consistency across function calls.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

…2_key

github-actions · 2026-01-15T03:19:34Z

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run pre-commit run --all-files in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀

coderabbitai · 2026-01-15T03:19:37Z

📝 Walkthrough

Walkthrough

The pull request includes formatting adjustments, quote style consistency updates, and argument parsing refactoring across example files and the builder module. Multi-line method calls are collapsed to single-line equivalents, exception print statements are updated to use double-quotes, and argument forwarding in the builder's get_tir method is refactored to use keyword-only propagation.

Changes

Cohort / File(s)	Summary
Example code formatting `examples/autodd/tilelang_buggy.py`, `examples/autodd/tilelang_minimized_expected.py`	Collapsed multi-line method calls to single-line format; removed extra blank lines; updated exception print statements from single to double-quote style for consistency
Builder argument refactoring `tilelang/language/eager/builder.py`	Refactored `get_tir` to use `_parse_phase1_key` for argument parsing; changed template construction to use keyword-only argument propagation (`*kwargs`) instead of positional and keyword arguments (`args, **kwargs`)

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

[AutoDD] Add Tilelang AutoDD to Reduce Buggy Program #1639 — Modifies the same example files and MatmulConfig methods with identical function refactoring patterns

Suggested labels

enhancement

Suggested reviewers

LeiWang1999
tzj-fxz

Poem

🐰 Hoppy changes hopping through the code,
Lines collapsed on the developer's road,
Quotes aligned in double-quote grace,
Arguments flow in their proper place! ✨

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 25.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately describes the main fix: addressing a bug in positional argument handling in the EagerJIT system, which is reflected in the key changes to argument parsing and forwarding in builder.py.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing touches

📝 Generate docstrings

📜 Recent review details

Configuration used: defaults

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 2d8d367 and 65b68f2.

📒 Files selected for processing (3)

examples/autodd/tilelang_buggy.py
examples/autodd/tilelang_minimized_expected.py
tilelang/language/eager/builder.py

🧰 Additional context used

🧬 Code graph analysis (1)

tilelang/language/eager/builder.py (1)

tilelang/jit/__init__.py (1)

get_tir (290-302)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)

GitHub Check: Test for Python 3.12 with Nightly-ROCm-7.1 (on self-hosted-amd)
GitHub Check: Test for Python 3.12 with Metal (on macos-latest)

🔇 Additional comments (5)

examples/autodd/tilelang_minimized_expected.py (1)

46-46: LGTM!

Minor formatting change to use double quotes for the f-string. No functional impact.

examples/autodd/tilelang_buggy.py (2)

75-79: LGTM!

Formatting improvements collapsing multi-line calls to single lines. No functional changes.

105-106: LGTM!

Whitespace formatting adjustment.

tilelang/language/eager/builder.py (2)

1062-1062: LGTM!

Correctly passes only non-tensor keyword arguments to _build_tir_template after tensor args have been extracted by _parse_phase1_key.

1067-1074: LGTM! Proper fix for positional argument handling.

The refactored get_tir method correctly:

Uses _parse_phase1_key to convert positional args to kwargs and extract tensor_args separately

Passes only non-tensor **kwargs to _build_tir_template

Passes **tensor_args, **kwargs to TirTemplate.get_tir, which expects keyword-only arguments

Since tensor_args are popped from kwargs in _parse_phase1_key (line 995), there's no risk of key collision when merging them back.

_{✏️ Tip: You can disable this entire section by setting review_details to false in your review settings.}

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

Fix argument unpacking in JITFunc and add debug print in _parse_phase…

65b68f2

…2_key

LeiWang1999 approved these changes Jan 15, 2026

View reviewed changes

LeiWang1999 merged commit 47aaf7b into tile-ai:main Jan 15, 2026
13 of 15 checks passed

LeiWang1999 mentioned this pull request Feb 3, 2026

Sync branch sync_0203: Features, Enhancements, and Bug Fixes #1780

Closed

3 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[EagerJIT] Fix bug in handling of positional arguments#1675

[EagerJIT] Fix bug in handling of positional arguments#1675
LeiWang1999 merged 1 commit intotile-ai:mainfrom
kurisu6912:eagerjit-param-bug

kurisu6912 commented Jan 15, 2026 •

edited

Loading

Uh oh!

github-actions bot commented Jan 15, 2026

Uh oh!

coderabbitai bot commented Jan 15, 2026 •

edited

Loading

Walkthrough

Changes

Estimated code review effort

Possibly related PRs

Suggested labels

Suggested reviewers

Poem

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

kurisu6912 commented Jan 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Uh oh!

github-actions bot commented Jan 15, 2026

Uh oh!

coderabbitai bot commented Jan 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Estimated code review effort

Possibly related PRs

Suggested labels

Suggested reviewers

Poem

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

kurisu6912 commented Jan 15, 2026 •

edited

Loading

coderabbitai bot commented Jan 15, 2026 •

edited

Loading